package com.ideaaedi.hot.clazz.update;

import com.ideaaedi.commonds.bash.BashUtil;
import com.ideaaedi.commonds.jar.JarUtil;
import com.ideaaedi.commonds.path.PathUtil;
import com.ideaaedi.hot.clazz.update.enums.TransferTypeEnum;
import com.ideaaedi.hot.clazz.update.exception.HotUpdateClassException;
import com.ideaaedi.hot.clazz.update.memory.MemoryTransferMap;
import com.ideaaedi.hot.common.HotUpdateExecutor;
import com.ideaaedi.hot.common.HotUpdateType;
import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * class热更新执行器
 *
 * @author JustryDeng
 * @since 1.0.0
 */
public class ClassHotUpdateInnerExecutor implements HotUpdateExecutor {
    
    public static final Logger log = LoggerFactory.getLogger(ClassHotUpdateInnerExecutor.class);
    
    /** jar文件后缀 */
    public static final String JAR_SUFFIX = ".jar";
    
    /** 同<version>2.2.0</version> */
    public static final String AGENT_JAR_VERSION = "2.2.0";
    
    /** class热更新的代理jar包包名 */
    public static final String AGENT_JAR_NAME = String.format("hot-update-class-%s.jar", AGENT_JAR_VERSION);
    
    /** 添加资源的方法 */
    private static final Method ADD_URL_METHOD;
    
    /** 类加载器 */
    private static final URLClassLoader CLASS_LOADER;
    
    static {
        try {
            ADD_URL_METHOD = URLClassLoader.class.getDeclaredMethod("addURL", URL.class);
            ADD_URL_METHOD.setAccessible(true);
            CLASS_LOADER = (URLClassLoader) ClassLoader.getSystemClassLoader();
        } catch (NoSuchMethodException e) {
            throw new HotUpdateClassException(e);
        }
    }
    
    /**
     * 传输类型
     */
    private final TransferTypeEnum transferType;
    
    public ClassHotUpdateInnerExecutor(TransferTypeEnum transferType) {
        this.transferType = transferType;
    }
    
    /**
     * 当要热更新的JVM进程是当前程序自身时，调用此方法可自动获取进程
     * <p></p>
     * @see ClassHotUpdateInnerExecutor#execute(String, String)
     */
    private void execute(String transferData) {
        // 获取当前JVM进程的进程id
        String pid = BashUtil.currPid();
        log.info("determine pid -> {}", pid);
        execute(pid, transferData);
    }
    
    /**
     * @see ClassHotUpdateInnerExecutor#execute(String, String, String)
     */
    private void execute(String pid, String transferData) {
        // 获取当前类所在的jar包的路径作为agentJarPath
        String agentJarPath = PathUtil.getProjectRootDir(ClassHotUpdateInnerExecutor.class);
        log.info("agentJarPath -> {}", agentJarPath);
        if (!agentJarPath.endsWith(AGENT_JAR_NAME)) {
            // 那么有可能AGENT_JAR_NAME作为lib包存在
            if (!agentJarPath.endsWith(JAR_SUFFIX)) {
                // 找不到代理类jar包
                throw new HotUpdateClassException(String.format("So, where is %s ? curr agentJarPath is [%s].", AGENT_JAR_NAME, agentJarPath));
            }
            String tmpDir = agentJarPath.substring(0, agentJarPath.lastIndexOf("/") + 1) + "__hot__temp__";
            List<String> list = new ArrayList<>();
            list.add("lib/" + AGENT_JAR_NAME);
            list.add("BOOT-INF/lib/" + AGENT_JAR_NAME);
            List<String> agentJarPathList = JarUtil.unJarWar(agentJarPath, tmpDir, true, list);
            if (agentJarPathList == null || agentJarPathList.size() == 0) {
                // 找不到代理类jar包
                throw new HotUpdateClassException(String.format("So, where is %s ?", AGENT_JAR_NAME));
            }
            agentJarPath = agentJarPathList.get(0);
            log.info("amend agentJarPath -> {}", agentJarPath);
        }
        execute(pid, agentJarPath, transferData);
    }
    
    /**
     * @see ClassHotUpdateInnerExecutor#doExecute(String, String, String)
     */
    private void execute(String pid, String agentJarPath, String transferData) {
        try {
            // 先确保tools.jar存在
            loadToolsJar();
            // 再往后执行逻辑
            doExecute(pid, agentJarPath, transferData);
        } catch (Exception e) {
            throw new HotUpdateClassException(e);
        }
    }
    
    /**
     * 执行热更新
     *
     * @param pid
     *         要被更新class的JVM进程id
     * @param agentJarPath
     *         {@link ClassHotUpdateInnerExecutor#AGENT_JAR_NAME}文件对应的全路径
     * @param transferData
     *            要传输到代理{@link ClassHotUpdateAgent#agentmain(String, Instrumentation)}的BASE64后的字符串。<br/>
     *            注：如果有多个BASE64后的字符串数据，那么多个BASE64字符串之间使用逗号分割。<br/>
     *            注：可使用{@link Base64.Encoder#encodeToString(byte[])}方法进行base64。
     */
    private void doExecute(String pid, String agentJarPath, String transferData) throws AgentLoadException, IOException,
            AgentInitializationException, AttachNotSupportedException {
        log.info("pid -> {}, agentJarPath -> {}, transferData -> {}", pid, agentJarPath, transferData);
        /*
         * 简单校验agentJarPath是否合法 （目标JVM进程 与 调用此方法时的进程应该是在同一台机器上的， 所以，
         * 校验目标JVM进程所在机器上是否存在agentJarPath文件，就相当于校验当前进程所在机器上是否存在agentJarPath文件
         */
        File agentJarFile = new File(agentJarPath);
        if (!agentJarFile.exists()) {
            throw new HotUpdateClassException("file agentJarPath non-exist.");
        }
        if (!agentJarFile.isFile()) {
            throw new HotUpdateClassException("file agentJarPath is not a file.");
        }
        if (!AGENT_JAR_NAME.equalsIgnoreCase(agentJarFile.getName())) {
            throw new HotUpdateClassException("AgentJarFile's name is not '" + AGENT_JAR_NAME + "'.");
        }
        
        VirtualMachine virtualMachine = VirtualMachine.attach(pid);
        // 指挥目标JVM加载对应agentJarPath位置的jar包作为代理，调用代理方法，其中transferData即为代理方法的入参
        virtualMachine.loadAgent(agentJarPath, wrapProtocol(transferData));
    }
    
    @Override
    public HotUpdateType type() {
        return HotUpdateType.CLASS;
    }
    
    /**
     * 添加协议
     *
     * @param transferData
     *            要传输的数据
     * @return  带有协议的数据
     */
    protected String wrapProtocol(String transferData) {
        return transferType.getProtocol() + transferData;
    }
    
    /**
     * 确保tools.jar存在
     */
    private void loadToolsJar() throws ClassNotFoundException, MalformedURLException, InvocationTargetException, IllegalAccessException {
        try {
            // VirtualMachine是否能加载到
            ClassHotUpdateInnerExecutor.class.getClassLoader().loadClass("com.sun.tools.attach.VirtualMachine");
            log.info("load VirtualMachine from classpath success.");
        } catch (ClassNotFoundException e) {
            log.warn("can not find tools.jar from classpath, try to load it from java.home.");
            String javaHome = System.getProperty("java.home");
            if (javaHome == null) {
                javaHome = System.getenv("JAVA_HOME");
            }
            if (javaHome == null) {
                log.warn("can not get java.home, can not load tools.jar.");
                throw e;
            }
    
            // java_home是jdk
            File jar = new File(javaHome + "/lib/tools.jar");
            if (!jar.exists()) {
                log.warn("can not find lib/tools.jar from java.home: {}", jar.getAbsolutePath());
                // java_home可能是jre
                jar = new File(javaHome + "/../lib/tools.jar");
                if (!jar.exists()) {
                    log.warn("can not find lib/tools.jar from java.home: {}", jar.getAbsolutePath());
                    /*
                     * 不排除这种目录结构：
                     *
                     *  ├─jdk1.8.0_181
                     *  │  │
                     *  │  └─lib
                     *  │      │──tools.jar
                     *  │      │
                     *  │
                     *  ├─jre1.8.0_181
                     */
                    jar = new File(javaHome + String.format("/../%s/lib/tools.jar",
                            new File(javaHome).getName().replace("jre", "jdk")));
                    if (!jar.exists()) {
                        log.warn("can not find lib/tools.jar from java.home: {}", jar.getAbsolutePath());
                        throw e;
                    }
                }
            }
            ADD_URL_METHOD.invoke(ClassHotUpdateInnerExecutor.class.getClassLoader(), jar.toURI().toURL());
            // 再次加载VirtualMachine是否能加载到
            ClassHotUpdateInnerExecutor.class.getClassLoader().loadClass("com.sun.tools.attach.VirtualMachine");
            log.info("reload VirtualMachine success.");
        }
    }
    
    /**
     * {@link TransferTypeEnum#MEMORY_CLASS_BYTES}策略下，对ClassHotUpdateExecutor的使用封装，使更方便使用
     *
     * @author JustryDeng
     * @since 2021/9/22 0:48:30
     */
    public static class MEMORY_CLASS_BYTES implements HotUpdateExecutor {
        
        /**
         * @see MEMORY_CLASS_BYTES#execute(String, List)
         */
        public static void execute(List<byte[]> classBytesList) {
            MEMORY_CLASS_BYTES.execute(null, classBytesList);
        }
        
        /**
         * 基于class字节码，走内存进行class热更新
         *
         * @param agentJarPath
         *            代理jar包的绝对路径
         * @param classBytesList
         *            要进行更新的class字节码文件的绝对路径集合
         */
        public static void execute(String agentJarPath, List<byte[]> classBytesList) {
            // 数据校验
            if (classBytesList == null || classBytesList.size() == 0) {
                throw new HotUpdateClassException("classBytesList cannot be empty.");
            }
            if (classBytesList.contains(null)) {
                throw new HotUpdateClassException("classBytesList cannot contains null.");
            }
            // 按TransferTypeEnum#MEMORY_CLASS_BYTES进行数据处理
            List<String> keyList = new ArrayList<>();
            String key;
            for (byte[] classBytes : classBytesList) {
                key = UUID.randomUUID().toString().replace("-", "");
                MemoryTransferMap.globalMemory().put(key, classBytes);
                keyList.add(key);
            }
            String transferData = keyList.stream()
                    .map(x -> Base64.getEncoder().encodeToString(x.getBytes(StandardCharsets.UTF_8)))
                    .collect(Collectors.joining(","));
    
            // 调用执行器执行热更新
            ClassHotUpdateInnerExecutor executor = new ClassHotUpdateInnerExecutor(TransferTypeEnum.MEMORY_CLASS_BYTES);
            if (StringUtils.isBlank(agentJarPath)) {
                executor.execute(transferData);
            } else {
                executor.execute(BashUtil.currPid(), agentJarPath, transferData);
            }
        }
    
        @Override
        public HotUpdateType type() {
            return HotUpdateType.CLASS;
        }
    }
    
    /**
     * {@link TransferTypeEnum#MEMORY_JAVA_BYTES}策略下，对ClassHotUpdateExecutor的使用封装，使更方便使用
     *
     * @author JustryDeng
     * @since 2021/9/22 0:48:30
     */
    public static class MEMORY_JAVA_BYTES implements HotUpdateExecutor {
        
        /**
         * @see MEMORY_JAVA_BYTES#execute(String, List)
         */
        public static void execute(List<byte[]> javaBytesList) {
            MEMORY_JAVA_BYTES.execute(null, javaBytesList);
        }
    
        /**
         * 基于java源码，走内存进行class热更新
         *
         * @param agentJarPath
         *            代理jar包的绝对路径
         * @param javaBytesList
         *            要进行更新的java源码文件的绝对路径集合
         */
        public static void execute(String agentJarPath, List<byte[]> javaBytesList) {
            // 数据校验
            if (javaBytesList == null || javaBytesList.size() == 0) {
                throw new HotUpdateClassException("javaBytesList cannot be empty.");
            }
            if (javaBytesList.contains(null)) {
                throw new HotUpdateClassException("javaBytesList cannot contains null.");
            }
            // 按TransferTypeEnum#MEMORY_JAVA_BYTES进行数据处理
            List<String> keyList = new ArrayList<>();
            String key;
            for (byte[] classBytes : javaBytesList) {
                key = UUID.randomUUID().toString().replace("-", "");
                MemoryTransferMap.globalMemory().put(key, classBytes);
                keyList.add(key);
            }
            String transferData = keyList.stream()
                    .map(x -> Base64.getEncoder().encodeToString(x.getBytes(StandardCharsets.UTF_8)))
                    .collect(Collectors.joining(","));
            
            // 调用执行器执行热更新
            ClassHotUpdateInnerExecutor executor = new ClassHotUpdateInnerExecutor(TransferTypeEnum.MEMORY_JAVA_BYTES);
            if (StringUtils.isBlank(agentJarPath)) {
                executor.execute(transferData);
            } else {
                executor.execute(BashUtil.currPid(), agentJarPath, transferData);
            }
        }
    
        @Override
        public HotUpdateType type() {
            return HotUpdateType.CLASS;
        }
    }
    
    /**
     * {@link TransferTypeEnum#BASE64_CLASS_PATHS}策略下，对ClassHotUpdateExecutor的使用封装，使更方便使用
     *
     * @author JustryDeng
     * @since 2021/9/22 0:48:30
     */
    public static class BASE64_CLASS_PATHS implements HotUpdateExecutor {
        
        /**
         * @see BASE64_CLASS_PATHS#execute(String, List)
         */
        public static void execute(List<String> classPathsList) {
            BASE64_CLASS_PATHS.execute(null, null, classPathsList);
        }
        
        /**
         * @see BASE64_CLASS_PATHS#execute(String, String, List)
         */
        public static void execute(String agentJarPath, List<String> classPathsList) {
            BASE64_CLASS_PATHS.execute(null, agentJarPath, classPathsList);
        }
        
        /**
         * 基于class字节码文件的class热更新
         *
         * @param pid
         *            要热更新的jvm的进程id
         * @param agentJarPath
         *            代理jar包的绝对路径
         * @param classPathsList
         *            要进行更新的class字节码文件的绝对路径集合
         */
        public static void execute(String pid, String agentJarPath, List<String> classPathsList) {
            // 数据校验
            if (classPathsList == null || classPathsList.size() == 0) {
                throw new HotUpdateClassException("classPathsList cannot be empty.");
            }
            if (classPathsList.contains(null)) {
                throw new HotUpdateClassException("classPathsList cannot contains null.");
            }
            // 按TransferTypeEnum#BASE64_CLASS_PATHS进行数据处理
            String transferData = classPathsList.stream()
                    .map(x -> Base64.getEncoder().encodeToString(x.getBytes(StandardCharsets.UTF_8)))
                    .collect(Collectors.joining(","));
            // 调用执行器执行热更新
            ClassHotUpdateInnerExecutor executor = new ClassHotUpdateInnerExecutor(TransferTypeEnum.BASE64_CLASS_PATHS);
            pid = StringUtils.isBlank(pid) ? BashUtil.currPid() : pid;
            if (StringUtils.isBlank(agentJarPath)) {
                executor.execute(pid, transferData);
            } else {
                executor.execute(pid, agentJarPath, transferData);
            }
        }
    
        @Override
        public HotUpdateType type() {
            return HotUpdateType.CLASS;
        }
    }
    
    /**
     * {@link TransferTypeEnum#BASE64_JAVA_PATHS}策略下，对ClassHotUpdateExecutor的使用封装，使更方便使用
     *
     * @author JustryDeng
     * @since 2021/9/22 0:48:30
     */
    public static class BASE64_JAVA_PATHS implements HotUpdateExecutor {
        
        /**
         * @see BASE64_JAVA_PATHS#execute(String, List)
         */
        public static void execute(List<String> javaPathsList) {
            BASE64_JAVA_PATHS.execute(null, null, javaPathsList);
        }
        
        /**
         * @see BASE64_JAVA_PATHS#execute(String, String, List)
         */
        public static void execute(String agentJarPath, List<String> javaPathsList) {
            BASE64_JAVA_PATHS.execute(null, agentJarPath, javaPathsList);
        }
        
        /**
         * 基于java源码文件的class热更新
         *
         * @param pid
         *            要热更新的jvm的进程id
         * @param agentJarPath
         *            代理jar包的绝对路径
         * @param javaPathsList
         *            要进行更新的java类文件的绝对路径集合
         */
        public static void execute(String pid, String agentJarPath, List<String> javaPathsList) {
            // 数据校验
            if (javaPathsList == null || javaPathsList.size() == 0) {
                throw new HotUpdateClassException("javaPathsList cannot be empty.");
            }
            if (javaPathsList.contains(null)) {
                throw new HotUpdateClassException("javaPathsList cannot contains null.");
            }
            // 按TransferTypeEnum#BASE64_CLASS_PATHS进行数据处理
            String transferData = javaPathsList.stream()
                    .map(x -> Base64.getEncoder().encodeToString(x.getBytes(StandardCharsets.UTF_8)))
                    .collect(Collectors.joining(","));
            // 调用执行器执行热更新
            ClassHotUpdateInnerExecutor executor = new ClassHotUpdateInnerExecutor(TransferTypeEnum.BASE64_JAVA_PATHS);
            pid = StringUtils.isBlank(pid) ? BashUtil.currPid() : pid;
            if (StringUtils.isBlank(agentJarPath)) {
                executor.execute(pid, transferData);
            } else {
                executor.execute(pid, agentJarPath, transferData);
            }
        }
    
        @Override
        public HotUpdateType type() {
            return HotUpdateType.CLASS;
        }
    }
}
